/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * License); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * AS IS BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*
 * Copyright (c) 2018, Open AI Lab
 * Author: chunyinglv@openailab.com
 */

#include <iostream>
#include <functional>
#include <cstring>
#include <algorithm>

#include "logger.hpp"
#include "node_ops.hpp"
#include "tensor_mem.hpp"
#include "graph.hpp"
#include "operator/rpn.hpp"
#include <math.h>

#ifndef MAX
#define MAX(a, b) (((a) > (b)) ? (a) : (b))
#endif
#ifndef MIN
#define MIN(a, b) (((a) < (b)) ? (a) : (b))
#endif
void _set(const int N, const float alpha, float* Y)
{
    for(int i = 0; i < N; ++i)
        Y[i] = alpha;
}

void _axpy(const int N, const float* X, float* Y)
{
    for(int i = 0; i < N; i++)
        Y[i] += X[i];
}
void _axpy_(const int N, const float* X, float* Y)
{
    for(int i = 0; i < N; i++)
        Y[i] -= X[i];
}
void _axpy_half(const int N, const float* X, float* Y)
{
    for(int i = 0; i < N; i++)
        Y[i] += 0.5 * X[i];
}
void _add_one(const int N, float* Y)
{
    for(int i = 0; i < N; i++)
        Y[i] += 1.f;
}
void _mul(const int N, const float* a, const float* b, float* y)
{
    for(int i = 0; i < N; i++)
        y[i] = a[i] * b[i];
}
void _add(const int N, const float* a, const float* b, float* y)
{
    for(int i = 0; i < N; i++)
        y[i] = a[i] + b[i];
}
void _exp(const int N, const float* a, float* y)
{
    for(int i = 0; i < N; i++)
        y[i] = exp(a[i]);
}

struct SBox
{
    float x0;
    float y0;
    float x1;
    float y1;
    float score;
    bool operator<(const SBox& tmp) const
    {
        return score < tmp.score;
    }
};
void proposal_local_anchor(int feat_height, int feat_width, int feat_stride, std::vector<Anchor>& anchors,
                           float* local_anchors)
{
    int length = std::max(feat_height, feat_width);
    int feat_size = feat_height * feat_width;
    int* map_m = new int[length];
    for(int i = 0; i < length; ++i)
    {
        map_m[i] = i * feat_stride;
    }
    float* shift_x = new float[feat_size];
    float* shift_y = new float[feat_size];
    for(int i = 0; i < feat_height; ++i)
    {
        for(int j = 0; j < feat_width; ++j)
        {
            shift_x[i * feat_width + j] = map_m[j];
            shift_y[i * feat_width + j] = map_m[i];
        }
    }
    int num_anchors = ( int )anchors.size();
    float* a = local_anchors;
    for(int i = 0; i < num_anchors; ++i)
    {
        _set(feat_size, anchors[i].x0, a + (i * 4 + 0) * feat_size);
        _set(feat_size, anchors[i].y0, a + (i * 4 + 1) * feat_size);
        _set(feat_size, anchors[i].x1, a + (i * 4 + 2) * feat_size);
        _set(feat_size, anchors[i].y1, a + (i * 4 + 3) * feat_size);
        _axpy(feat_size, shift_x, a + (i * 4 + 0) * feat_size);
        _axpy(feat_size, shift_x, a + (i * 4 + 2) * feat_size);
        _axpy(feat_size, shift_y, a + (i * 4 + 1) * feat_size);
        _axpy(feat_size, shift_y, a + (i * 4 + 3) * feat_size);
    }
    delete[] map_m;
    delete[] shift_x;
    delete[] shift_y;
}

void bbox_tranform_inv(float* m_box, float* local_anchors, int height, int width, int channel, int num_anchors)
{
    int step = height * width;
    float* a = m_box;
    float* b = local_anchors;
    int c_4 = channel / 4;
    for(int i = 0; i < c_4; ++i)
    {
        _axpy_(2 * step, b + (i * 4 + 0) * step, b + (i * 4 + 2) * step);
        _add_one(2 * step, b + (i * 4 + 2) * step);
        _axpy_half(2 * step, b + (i * 4 + 2) * step, b + (i * 4 + 0) * step);

        _mul(2 * step, b + (i * 4 + 2) * step, a + (i * 4 + 0) * step, a + (i * 4 + 0) * step);
        _add(2 * step, b + (i * 4 + 0) * step, a + (i * 4 + 0) * step, a + (i * 4 + 0) * step);

        _exp(2 * step, a + (i * 4 + 2) * step, a + (i * 4 + 2) * step);
        _mul(2 * step, b + (i * 4 + 2) * step, a + (i * 4 + 2) * step, a + (i * 4 + 2) * step);
    }
}

void filter_boxs(std::vector<SBox>& boxes, float* box, float* score, int min_size, int src_scale, int src_w, int src_h,
                 int feat_w, int feat_h, int num_anchors, int feat_c)
{
    float local_minsize = min_size * src_scale;
    boxes.clear();

    int feat_c_ = feat_c / 4;
    int one_step = feat_h * feat_w;
    int step = 4 * one_step;
    int offset_w, offset_h, offset_x, offset_y, offset_s;

    for(int h = 0; h < feat_h; ++h)
    {
        for(int w = 0; w < feat_w; ++w)
        {
            offset_x = h * feat_w + w;
            offset_y = offset_x + one_step;
            offset_w = offset_y + one_step;
            offset_h = offset_w + one_step;
            offset_s = one_step * num_anchors + offset_x;
            for(int c = 0; c < feat_c_; ++c)
            {
                float width = box[offset_w];
                float height = box[offset_h];
                if((width >= local_minsize) & (height >= local_minsize))
                {
                    SBox tmp;
                    tmp.x0 = box[offset_x] - 0.5 * width;
                    tmp.y0 = box[offset_y] - 0.5 * height;
                    tmp.x1 = box[offset_x] + 0.5 * width;
                    tmp.y1 = box[offset_y] + 0.5 * height;
                    tmp.x0 = MIN(MAX(tmp.x0, 0), src_w);
                    tmp.y0 = MIN(MAX(tmp.y0, 0), src_h);
                    tmp.x1 = MIN(MAX(tmp.x1, 0), src_w);
                    tmp.y1 = MIN(MAX(tmp.y1, 0), src_h);
                    tmp.score = score[offset_s];
                    boxes.push_back(tmp);
                }
                offset_x += step;
                offset_y += step;
                offset_w += step;
                offset_h += step;
                offset_s += one_step;
            }
        }
    }
}

void nms_rpn(std::vector<SBox>& input_boxes, float nms_thresh)
{
    std::vector<float> vArea(input_boxes.size());
    for(int i = 0; i < ( int )input_boxes.size(); ++i)
    {
        vArea[i] =
            (input_boxes.at(i).x1 - input_boxes.at(i).x0 + 1) * (input_boxes.at(i).y1 - input_boxes.at(i).y0 + 1);
    }
    for(int i = 0; i < ( int )input_boxes.size(); ++i)
    {
        for(int j = i + 1; j < ( int )input_boxes.size();)
        {
            float xx1 = std::max(input_boxes[i].x0, input_boxes[j].x0);
            float yy1 = std::max(input_boxes[i].y0, input_boxes[j].y0);
            float xx2 = std::min(input_boxes[i].x1, input_boxes[j].x1);
            float yy2 = std::min(input_boxes[i].y1, input_boxes[j].y1);
            float w = std::max(float(0), xx2 - xx1 + 1);
            float h = std::max(float(0), yy2 - yy1 + 1);
            float inter = w * h;
            float ovr = inter / (vArea[i] + vArea[j] - inter);
            if(ovr >= nms_thresh)
            {
                input_boxes.erase(input_boxes.begin() + j);
                vArea.erase(vArea.begin() + j);
            }
            else
            {
                j++;
            }
        }
    }
}
namespace TEngine {

namespace RPNImpl {

struct RPNOps : public NodeOps
{
    bool Run(Node* node)
    {
        const Tensor* score_tensor = node->GetInputTensor(0);
        const Tensor* featmap_tensor = node->GetInputTensor(1);
        const Tensor* info_tensor = node->GetInputTensor(2);
        Tensor* output_tensor = node->GetOutputTensor(0);
        TShape& out_shape = output_tensor->GetShape();

        float* output = ( float* )get_tensor_mem(output_tensor);
        const float* im_info = ( float* )get_tensor_mem(info_tensor);
        const float* m_score_ = ( float* )get_tensor_mem(score_tensor);
        const float* m_box_ = ( float* )get_tensor_mem(featmap_tensor);

        const TShape& featmap_shape = featmap_tensor->GetShape();
        const TShape& score_shape = score_tensor->GetShape();
        const int feat_height = featmap_shape.GetH();
        const int feat_width = featmap_shape.GetW();
        const int feat_channel = featmap_shape.GetC();
        const int score_channel = score_shape.GetC();
        const int feat_size = feat_height * feat_width;
        int src_height_ = im_info[0];
        int src_width_ = im_info[1];
        int src_scale_ = im_info[2];

        RPN* RPN_op = dynamic_cast<RPN*>(node->GetOp());
        RPNParam* param_ = RPN_op->GetParam();
        int feat_stride = param_->feat_stride;
        int num_anchors = ( int )param_->anchors_.size();

        // local_anchors (1, anchors_nums_ * 4, map_height_, map_width_);
        float* local_anchors = new float[num_anchors * 4 * feat_size];
        proposal_local_anchor(feat_height, feat_width, feat_stride, param_->anchors_, local_anchors);

        float* m_box = new float[feat_channel * feat_size];
        for(int i = 0; i < feat_channel * feat_size; i++)
            m_box[i] = m_box_[i];
        bbox_tranform_inv(m_box, local_anchors, feat_height, feat_width, feat_channel, num_anchors);

        delete[] local_anchors;
        std::vector<SBox> boxes;
        float* m_score = new float[score_channel * feat_size];
        for(int i = 0; i < score_channel * feat_size; i++)
            m_score[i] = m_score_[i];
        filter_boxs(boxes, m_box, m_score, param_->min_size, src_scale_, src_width_, src_height_, feat_width,
                    feat_height, num_anchors, feat_channel);
        delete[] m_box;
        delete[] m_score;

        std::sort(boxes.rbegin(), boxes.rend());

        if(param_->per_nms_topn > 0)
        {
            int tmp = MIN(param_->per_nms_topn, ( int )boxes.size());
            boxes.erase(boxes.begin() + tmp, boxes.end());
        }
        nms_rpn(boxes, param_->nms_thresh);

        if(param_->post_nms_topn > 0)
        {
            int tmp = MIN(param_->post_nms_topn, ( int )boxes.size());

            boxes.erase(boxes.begin() + tmp, boxes.end());
        }
        // inder shape [default batch=1]
        int num_box = boxes.size();
        std::vector<int> outdim = {1, num_box, 4, 1};
        out_shape.SetDim(outdim);

        float* out_data = output;
        // std::cout<<"num_box "<<num_box<<"\n";
        for(int i = 0; i < num_box; i++)
        {
            const SBox& r = boxes[i];
            float* outptr = out_data + i * 4;
            outptr[0] = r.x0;
            outptr[1] = r.y0;
            outptr[2] = r.x1;
            outptr[3] = r.y1;
        }
        return true;
    }
};

NodeOps* SelectFunc(const CPUInfo* cpu_info, Node* node)
{
    Tensor* input = node->GetInputTensor(0);
    const int data_type = input->GetDataType();
    const ExecAttr* exec_attr = any_cast<const ExecAttr*>(node->GetAttr(ATTR_EXEC_ATTR));
    if(data_type != TENGINE_DT_FP32 || exec_attr->graph_layout != TENGINE_LAYOUT_NCHW)
        return nullptr;

    RPNOps* ops = new RPNOps();

    return ops;
}

}    // namespace RPNImpl

using namespace RPNImpl;

void RegisterRPNNodeExec(void)
{
    NodeOpsRegistryManager::RegisterOPImplementor("common", "RPN", RPNImpl::SelectFunc, 1000);
}

}    // namespace TEngine
